#!/usr/bin/env python
#
# Copyright 2010 Sebible Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
"""A Backend class to wrapp python-ldap to build a working backend for our QuerySet"""


import ldap
import ldap.schema
import ldapurl
import ldap.sasl

from django.utils.encoding import smart_str, smart_unicode
from django.db.backends import BaseDatabaseWrapper, BaseDatabaseFeatures, BaseDatabaseOperations
from django.conf import settings

from django_ldap import exceptions

import logging
logger = logging.getLogger('django_ldap.backend')

class DatabaseOperations(BaseDatabaseOperations):
    """docstring for DatabaseOperations"""
    def quote_name(self, name):
        """
        Implement the quote_name function. Since we don't have to quote name in ldap functions,
        we simply returns the name without quote.
        """
        return name


class DatabaseWrapper(BaseDatabaseWrapper):

    features = BaseDatabaseFeatures()
    ops = DatabaseOperations()
    operators = {
        'exact': '=%s',
        'iexact': '~=%s',
        'contains': '=*%s*',
        'icontains': '=*%s*',
        'regex': 'REGEXP BINARY %s',
        'iregex': 'REGEXP %s',
        'gt': '>%s',
        'gte': '>=%s',
        'lt': '<%s',
        'lte': '<=%s',
        'startswith': '=%s*',
        'endswith': '=*%s',
        'istartswith': '=%s*',
        'iendswith': '=*%s',
    }

    def __init__(self, **kwargs):
        """
        @todo: set self.server_version as the ldap server version
        """
        super(DatabaseWrapper, self).__init__(**kwargs)
        self.server_version = None

    def _valid_connection(self):
        """
        Since ldap connection don't have the connection.ping() function to verify the 
        connection, so we just simply returns true if self.connection is not None.
        """
        if self.connection:
            return True
        else:
            return False

    def _cursor(self, settings):
        """
        The function logic is almost the same as django.db.backends.DatabaseWrapper, but
        the ldap connection don't have the cursor() function. So we just simply returns 
        the connection as its cursor.
        """
        if not self._valid_connection():
            kwargs = {}
            if settings.LDAP_SERVER:
                kwargs['server'] = settings.LDAP_SERVER
            if settings.LDAP_BASE:
                kwargs['base'] = settings.LDAP_BASE
            if settings.LDAP_BIND_DN:
                kwargs['bind_dn'] = settings.LDAP_BIND_DN
            if settings.LDAP_BIND_PWD:
                kwargs['bind_pwd'] = settings.LDAP_BIND_PWD
            kwargs.update(self.options)
            self.connection = Backend(**kwargs)
        return self.connection


class Backend(object):
    def __init__( self, server='localhost', port=389, base='', use_ssl=0, bind_dn='', bind_pwd='', read_only=0, use_ad=False, **kwargs):
        """ Create a new LDAPDelegate instance """
        self.server = server
        self.port = port
        self.base = base
        self.use_ssl = use_ssl
        self.bind_dn = bind_dn
        self.bind_pwd = bind_pwd
        self.read_only = read_only
        self.use_ad = use_ad
        self.connection = None

    def get_ldapuri(self):
        if self.use_ssl:
            protocol = 'ldaps'
        else:
            protocol = 'ldap'
        uri = '%s://%s:%s' % (protocol, self.server, self.port)
        return uri

    def get_connection(self):
        if self.connection:
            return self.connection
        self.connection = self._connect(self.get_ldapuri(), self.bind_dn, self.bind_pwd)
        return self.connection

    def handle_referral(self, exception):
        """ Handle a referral specified in a exception """
        payload = exception.args[0]
        info = payload.get('info')
        ldap_url = info[info.find('ldap'):]

        if ldapurl.isLDAPUrl(ldap_url):
            parsed_url = ldapurl.LDAPUrl(ldap_url)
            conn_str = '%s://%s' % ( parsed_url.urlscheme or 'ldap'
                                   , parsed_url.hostport
                                   )

            return self._connect(conn_str, self.bind_dn, self.bind_pwd, 1)

        else:
            raise exceptions.LdapBackendException('Bad referral "%s"' % str(exception))

    def _connect(self, uri, bind_dn, bind_pwd, handle_referral=0):
        try:
            connection_type = ldap.ldapobject.SmartLDAPObject
        except AttributeError:
            try:
                connection_type = ldap.ldapobject.ReconnectLDAPObject
            except AttributeError:
                connection_type = ldap.ldapobject.SimpleLDAPObject
        try:
            connection = connection_type(uri,start_tls=self.use_ssl)
        except ldap.SERVER_DOWN:
            raise exceptions.LdapBackendException('Cannot connect to LDAP server')
        # Set the protocol version - version 3 is preferred
        try:
            connection.set_option(ldap.OPT_PROTOCOL_VERSION, ldap.VERSION3)
        except ldap.LDAPError: # Invalid protocol version, fall back safely
            connection.set_option(ldap.OPT_PROTOCOL_VERSION, ldap.VERSION2)

        # Deny auto-chasing of referrals to be safe, we handle them instead
        try:
            connection.set_option(ldap.OPT_REFERRALS, 0)
        except ldap.LDAPError: # Cannot set referrals, so do nothing
            pass

        # Now bind with the credentials given. Let exceptions propagate out.
        try:
            if not self.use_ad:
                connection.simple_bind_s(bind_dn, bind_pwd)
            else:
                sasl_auth = ldap.sasl.sasl(
                  {
                    ldap.sasl.CB_AUTHNAME:bind_dn,
                    ldap.sasl.CB_PASS    :bind_pwd,
                  },
                    'DIGEST-MD5'
                )
                connection.sasl_interactive_bind_s("", sasl_auth)
        except ldap.INVALID_CREDENTIALS:
            raise exceptions.LdapBackendException('Invalid authentication credentials')
        except:
            import sys
            tp,v,tb = sys.exc_info()
            raise exceptions.LdapBackendException('[%s][%s]%s' % (tb.tb_lineno, tp, v))

        try: # XXX Not sure if needed.
            connection.manage_dsa_it(handle_referral)
        except:
            pass

        return connection

    def search( self
              , filter='(objectClass=*)'
              , base = None
              , scope = None
              , attrs=[]
              ):
        """ The main search engine """
        if not base:
            base = self.base
        if not scope:
            scope = ldap.SCOPE_SUBTREE
        result = []
        filter = smart_str(filter)
        base = smart_str(base)
        connection = self.get_connection()
        try:
            try:
                res = connection.search_s(base, scope, filter, attrs)
            except ldap.PARTIAL_RESULTS:
                res_type, res = connection.result(all=0)
            except ldap.REFERRAL, e:
                connection = self.handle_referral(e)

                try:
                    res = connection.search_s(base, scope, filter, attrs)
                except ldap.PARTIAL_RESULTS:
                    res_type, res = connection.result(all=0)

            for rec_dn, rec_dict in res:
                # When used against Active Directory, "rec_dict" may not be
                # be a dictionary in some cases (instead, it can be a list)
                # An example of a useless "res" entry that can be ignored
                # from AD is
                # (None, ['ldap://ForestDnsZones.PORTAL.LOCAL/DC=ForestDnsZones,DC=PORTAL,DC=LOCAL'])
                # This appears to be some sort of internal referral, but
                # we can't handle it, so we need to skip over it.
                try:
                    items =  rec_dict.items()
                except AttributeError:
                    # 'items' not found on rec_dict
                    continue

                for key, value in items:
                    if not isinstance(value, str):
                        try:
                            for i in range(len(value)):
                                value[i] = smart_unicode(value[i])
                        except:
                            pass

                rec_dict['dn'] = [smart_unicode(rec_dn)]

                result.append(rec_dict)

        except ldap.NO_SUCH_OBJECT:
            raise exceptions.LdapBackendNoSuchObjectException('Cannot find %s under %s' % (filter, base))

        except ldap.SIZELIMIT_EXCEEDED:
            raise exceptions.LdapBackendException('Too many results for this query')

        except:
            import sys
            tp,v,tb = sys.exc_info()
            raise exceptions.LdapBackendException('[%s][%s]%s' % (tb.tb_lineno, tp, v))

        return result

    def insert(self, dn, attrs=None):
        """ Insert a new record """
        if self.read_only:
            raise exceptions.LdapBackendException('Running in read-only mode, insertion is disabled')

        dn = smart_str(dn)
        attribute_list = []
        attrs = attrs and attrs or {}

        for attr_key, attr_val in attrs.items():
            if isinstance(attr_val, str) or isinstance(attr_val, unicode):
                attr_val = [x.strip() for x in attr_val.split(';')]

            if attr_val != [''] and attr_val != None:
                attr_val = [smart_str(value, strings_only=True) for value in attr_val if value is not None]
                if len(attr_val) > 0:
                    attribute_list.append((attr_key, attr_val))
        connection = self.get_connection()

        try:
            connection.add_s(dn, attribute_list)
        except ldap.INVALID_CREDENTIALS, e:
            e_name = e.__class__.__name__
            raise exceptions.LdapBackendException('%s No permission to insert "%s"' % (e_name, dn))
        except ldap.ALREADY_EXISTS, e:
            e_name = e.__class__.__name__
            raise exceptions.LdapBackendException('%s Record with dn "%s" already exists' % (e_name, dn))
        except ldap.REFERRAL, e:
            try:
                connection = self.handle_referral(e)
                connection.add_s(dn, attribute_list)
            except ldap.INVALID_CREDENTIALS:
                e_name = e.__class__.__name__
                raise exceptions.LdapBackendException('%s No permission to insert "%s"' % (e_name, dn))
            except Exception, e:
                e_name = e.__class__.__name__
                raise exceptions.LdapBackendException('%s LDAPDelegate.insert: %s' % (e_name, str(e)))
        except Exception, e:
            e_name = e.__class__.__name__
            raise exceptions.LdapBackendException('%s django_ldap.backend.Backend.insert: %s' % (e_name, str(e)))

    def delete(self, dn):
        """ Delete a record """
        if self.read_only:
            raise exceptions.LdapBackendException('Running in read-only mode, insertion is disabled')

        connection = self.get_connection()
        try:
            connection.delete_s(dn)
        except ldap.INVALID_CREDENTIALS:
            raise exceptions.LdapBackendException('No permission to delete "%s"' % dn)
        except ldap.REFERRAL, e:
            try:
                connection = self.handle_referral(e)
                connection.delete_s(dn)
            except ldap.INVALID_CREDENTIALS:
                raise exceptions.LdapBackendException('No permission to delete "%s"' % dn)
            except Exception, e:
                raise exceptions.LdapBackendException('django_ldap.backend.Backend.delete: %s' % str(e))
        except Exception, e:
            raise exceptions.LdapBackendException('django_ldap.backend.Backend.delete: %s' % str(e))

    def save_history(self, dn, mod_list):
        pass

    def modify(self, dn, attrs=None, mod_type=None):
        """ Modify a record """
        if self.read_only:
            raise exceptions.LdapBackendException('Running in read-only mode, modification is disabled')

        utf8_dn = smart_str(dn)

        res = self.search(base=utf8_dn, scope=ldap.SCOPE_BASE)
        attrs = attrs and attrs or {}


        if len(res) == 0:
            raise exceptions.LdapBackendException('django_ldap.backend.Backend.modify: Cannot find dn "%s"' % dn)

        cur_rec = res[0]
        mod_list = []


        for key, values in attrs.items():
            if values is None:
                values = []
            values = [smart_str(value, strings_only=True) for value in values if value is not None]

            if mod_type is None:
                if map(smart_str, cur_rec.get(key, [''])) != values and values != [''] and len(values) > 0:
                    mod_list.append((ldap.MOD_REPLACE, key, values))
                elif cur_rec.has_key(key) and (values == [''] or values == [] or values is None):
                    mod_list.append((ldap.MOD_DELETE, key, None))
            else:
                mod_list.append((mod_type, key, values))
        connection = self.get_connection()
        try:
            old_dn_exploded = ldap.explode_dn(utf8_dn)
            rdn, rdn_value = old_dn_exploded[0].split('=')
            new_rdn_value = attrs.get(rdn,[''])
            if new_rdn_value and rdn_value!=new_rdn_value[0]:
                new_utf8_rdn = smart_str('%s=%s' % (rdn, new_rdn_value[0]))
                connection.modrdn_s(utf8_dn, new_utf8_rdn)
                old_dn_exploded[0] = new_utf8_rdn
                utf8_dn = ','.join(old_dn_exploded)
            connection.modify_s(utf8_dn, mod_list)
        except ldap.INVALID_CREDENTIALS, e:
            e_name = e.__class__.__name__
            raise exceptions.LdapBackendException('%s No permission to modify "%s"' % (e_name, dn))

        except ldap.REFERRAL, e:
            try:
                connection = self.handle_referral(e)
                connection.modify_s(dn, mod_list)
            except ldap.INVALID_CREDENTIALS, e:
                e_name = e.__class__.__name__
                raise exceptions.LdapBackendException('%s No permission to modify "%s"' % (e_name, dn))
            except Exception, e:
                e_name = e.__class__.__name__
                raise exceptions.LdapBackendException('%s django_ldap.backend.Backend.modify: %s' % (e_name, str(e)))

        except Exception, e:
            e_name = e.__class__.__name__
            raise exceptions.LdapBackendException('%s django_ldap.backend.Backend.modify: %s' % (e_name, str(e)))

    def get_all_attributes(self, name):
        if not hasattr(self, '_schema_cache'):
            subschemasubentry_dn, self._schema_cache = ldap.schema.urlfetch(self.get_ldapuri())
            if subschemasubentry_dn is None:
                raise exceptions.LdapBackendException('No sub schema sub entry found!')
        obj = self._schema_cache.get_obj(ldap.schema.ObjectClass, name)
        attrs = set(obj.must)
        attrs.update(obj.may)
        if hasattr(obj,'sup'):
            for sup in obj.sup:
                attrs.update(self.get_all_attributes(sup))
        return attrs

connection = DatabaseWrapper(**settings.LDAP_OPTIONS)